"""Prompt template management system."""

from typing import Dict, List, Any, Optional
from datetime import datetime
import logging
import jinja2
from motor.motor_asyncio import AsyncIOMotorClient
from bson import ObjectId

logger = logging.getLogger(__name__)


class PromptManager:
    """Manager for prompt templates with CRUD operations and versioning."""
    
    def __init__(self, mongodb_client: AsyncIOMotorClient):
        """Initialize prompt manager.
        
        Args:
            mongodb_client: MongoDB client for template storage
        """
        self.db = mongodb_client.ai_writing
        self.templates = self.db.prompt_templates
        self.jinja_env = jinja2.Environment(
            autoescape=True,
            trim_blocks=True,
            lstrip_blocks=True
        )
    
    async def create_prompt(
        self,
        name: str,
        category: str,
        system_prompt: str,
        user_prompt: str,
        variables: List[str],
        examples: Optional[List[Dict[str, str]]] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        """Create a new prompt template.
        
        Args:
            name: Template name
            category: Template category for organization
            system_prompt: System message template
            user_prompt: User message template
            variables: List of variable names used in templates
            examples: Optional examples showing template usage
            metadata: Optional additional metadata
            
        Returns:
            Created template ID
        """
        template = {
            "name": name,
            "category": category,
            "system_prompt": system_prompt,
            "user_prompt": user_prompt,
            "variables": variables,
            "examples": examples or [],
            "metadata": metadata or {},
            "version": 1,
            "created_at": datetime.utcnow(),
            "updated_at": datetime.utcnow(),
            "is_active": True
        }
        
        try:
            result = await self.templates.insert_one(template)
            logger.info(f"Created prompt template: {name} (ID: {result.inserted_id})")
            return str(result.inserted_id)
        except Exception as e:
            logger.error(f"Error creating prompt template: {e}")
            raise
    
    async def update_prompt(
        self,
        template_id: str,
        updates: Dict[str, Any],
        create_version: bool = True
    ) -> bool:
        """Update an existing prompt template.
        
        Args:
            template_id: ID of template to update
            updates: Dictionary of fields to update
            create_version: Whether to create a new version
            
        Returns:
            True if update successful
        """
        try:
            # Get current template
            current = await self.templates.find_one({"_id": ObjectId(template_id)})
            if not current:
                logger.error(f"Template {template_id} not found")
                return False
            
            # Create version history if requested
            if create_version:
                # Archive current version
                current["is_active"] = False
                current["archived_at"] = datetime.utcnow()
                await self.templates.insert_one(current)
                
                # Increment version number
                updates["version"] = current.get("version", 1) + 1
            
            # Update timestamp
            updates["updated_at"] = datetime.utcnow()
            
            # Perform update
            result = await self.templates.update_one(
                {"_id": ObjectId(template_id), "is_active": True},
                {"$set": updates}
            )
            
            return result.modified_count > 0
        except Exception as e:
            logger.error(f"Error updating prompt template: {e}")
            raise
    
    async def get_prompt(
        self,
        template_id: str,
        version: Optional[int] = None
    ) -> Optional[Dict[str, Any]]:
        """Get a prompt template by ID.
        
        Args:
            template_id: ID of template to retrieve
            version: Optional specific version to retrieve
            
        Returns:
            Template dictionary or None if not found
        """
        try:
            query = {"_id": ObjectId(template_id)}
            if version:
                query["version"] = version
            else:
                query["is_active"] = True
            
            template = await self.templates.find_one(query)
            if template:
                template["_id"] = str(template["_id"])
            return template
        except Exception as e:
            logger.error(f"Error getting prompt template: {e}")
            return None
    
    async def list_prompts(
        self,
        category: Optional[str] = None,
        limit: int = 100,
        skip: int = 0
    ) -> List[Dict[str, Any]]:
        """List prompt templates.
        
        Args:
            category: Optional category filter
            limit: Maximum number of templates to return
            skip: Number of templates to skip
            
        Returns:
            List of template dictionaries
        """
        try:
            query = {"is_active": True}
            if category:
                query["category"] = category
            
            cursor = self.templates.find(query).skip(skip).limit(limit)
            templates = await cursor.to_list(length=limit)
            
            # Convert ObjectId to string
            for template in templates:
                template["_id"] = str(template["_id"])
            
            return templates
        except Exception as e:
            logger.error(f"Error listing prompt templates: {e}")
            return []
    
    async def delete_prompt(self, template_id: str, soft_delete: bool = True) -> bool:
        """Delete a prompt template.
        
        Args:
            template_id: ID of template to delete
            soft_delete: If True, mark as inactive; if False, permanently delete
            
        Returns:
            True if deletion successful
        """
        try:
            if soft_delete:
                # Soft delete - mark as inactive
                result = await self.templates.update_one(
                    {"_id": ObjectId(template_id)},
                    {
                        "$set": {
                            "is_active": False,
                            "deleted_at": datetime.utcnow()
                        }
                    }
                )
            else:
                # Hard delete - remove from database
                result = await self.templates.delete_one(
                    {"_id": ObjectId(template_id)}
                )
            
            return result.modified_count > 0 if soft_delete else result.deleted_count > 0
        except Exception as e:
            logger.error(f"Error deleting prompt template: {e}")
            raise
    
    async def render_prompt(
        self,
        template_id: str,
        variables: Dict[str, Any]
    ) -> Dict[str, str]:
        """Render a prompt template with variables.
        
        Args:
            template_id: ID of template to render
            variables: Variables to inject into template
            
        Returns:
            Dictionary with rendered system and user prompts
        """
        template = await self.get_prompt(template_id)
        if not template:
            raise ValueError(f"Template {template_id} not found")
        
        try:
            # Create Jinja2 templates
            system_tmpl = self.jinja_env.from_string(template["system_prompt"])
            user_tmpl = self.jinja_env.from_string(template["user_prompt"])
            
            # Render templates
            rendered = {
                "system": system_tmpl.render(**variables),
                "user": user_tmpl.render(**variables)
            }
            
            return rendered
        except jinja2.TemplateError as e:
            logger.error(f"Error rendering template: {e}")
            raise ValueError(f"Template rendering failed: {e}")
    
    async def get_templates_by_category(self, category: str) -> List[Dict[str, Any]]:
        """Get all templates in a specific category.
        
        Args:
            category: Category name
            
        Returns:
            List of template dictionaries
        """
        return await self.list_prompts(category=category)
    
    async def get_template_versions(self, name: str) -> List[Dict[str, Any]]:
        """Get all versions of a template by name.
        
        Args:
            name: Template name
            
        Returns:
            List of template versions
        """
        try:
            cursor = self.templates.find({"name": name}).sort("version", -1)
            templates = await cursor.to_list(length=100)
            
            for template in templates:
                template["_id"] = str(template["_id"])
            
            return templates
        except Exception as e:
            logger.error(f"Error getting template versions: {e}")
            return []
    
    async def validate_template(
        self,
        system_prompt: str,
        user_prompt: str,
        variables: List[str]
    ) -> Dict[str, Any]:
        """Validate a template before saving.
        
        Args:
            system_prompt: System message template
            user_prompt: User message template
            variables: List of expected variables
            
        Returns:
            Validation result dictionary
        """
        result = {
            "valid": True,
            "errors": [],
            "warnings": []
        }
        
        try:
            # Parse templates to check syntax
            system_tmpl = self.jinja_env.from_string(system_prompt)
            user_tmpl = self.jinja_env.from_string(user_prompt)
            
            # Get variables from templates
            system_vars = set(system_tmpl.module.__dict__.get('_module', {}).keys())
            user_vars = set(user_tmpl.module.__dict__.get('_module', {}).keys())
            template_vars = system_vars.union(user_vars)
            
            # Check for missing variables
            expected_vars = set(variables)
            missing = expected_vars - template_vars
            if missing:
                result["warnings"].append(f"Variables not used in template: {missing}")
            
            # Check for undefined variables
            extra = template_vars - expected_vars
            if extra:
                result["errors"].append(f"Undefined variables in template: {extra}")
                result["valid"] = False
                
        except jinja2.TemplateError as e:
            result["errors"].append(f"Template syntax error: {e}")
            result["valid"] = False
        
        return result